#!/usr/bin/env python

import xlh
import sys, getopt, os.path
from math import log10, ceil, floor
from time import time, ctime
from scipy.stats import percentileofscore, scoreatpercentile
import numpy as np

import pdb


def correct(value, Max = 100):
	# apparently there's a bug in scipy.stats.percentileofscore...
	if value > Max:
		value -= 100
	return round(value, 2)

def combinedPV(xlinteraction):
	p = 1.0
	for m in xlinteraction.Xlinks():
		p *=  m.pValue()
	try:
		return -log10(p)
	except OverflowError:
		return 'Inf'
		
def printMatchingIons(xlink, tolerance, unit):
	spectrum = xlink.spectrum()
#	spectrum = xlink.spectrum().filterOnIntensity(0.01 * xlink.spectrum().mostIntense().peakIntensities()[0])
	NPeptide = xlink.NPeptide()
	CPeptide = xlink.CPeptide()
	
	NSpool = []
	CSpool = []
	
#	print NPeptide, xlink.xStart(), NPeptide.sequence(), NPeptide.sequence()[xlink.xStart()]
#	print CPeptide, xlink.xEnd(), CPeptide.sequence(), CPeptide.sequence()[xlink.xEnd()]
	#sp.closeTo(mz, tol, u).peakMasses()
	# for now print only b and y ions of simple peptides

	for s in range(xlink.xStart() - 1):
		index = s + 1
		mz = NPeptide.bMass(index)
		found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
		row = 'b'+str(index) + ' ' +  NPeptide.sequence()[s] + ' ' + str(NPeptide.bMass(index)) + ' ' +  str(found)
		NSpool.append(row)
	for s in range(xlink.xStart(), NPeptide.length()):
		index = NPeptide.length() - s
		mz = NPeptide.yMass(index)
		found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
		row = 'y'+str(index) + ' ' +  NPeptide.sequence()[s] + ' ' +  str(NPeptide.yMass(index)) + ' ' + str(found)
		NSpool.append(row)

	for s in range(xlink.xEnd() - 1):
		index = s + 1
		mz = CPeptide.bMass(index)
		found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
		row = 'b'+str(index) + ' ' + CPeptide.sequence()[s] + ' ' + str(CPeptide.bMass(index)) + ' ' + str(found)
		CSpool.append(row)
	for s in range(xlink.xEnd(), CPeptide.length()):
		index = CPeptide.length() - s
		mz = CPeptide.yMass(index)
		found = spectrum.closeTo(mz, tolerance, unit).closestTo(mz)
		row = 'y'+str(index) + ' ' + CPeptide.sequence()[s] + ' ' + str(CPeptide.yMass(index)) + ' ' + str(found)
		CSpool.append(row)
		
	NS = max(map(lambda x: len(x), NSpool))

	print NPeptide.sequence(), ' ' * (NS - NPeptide.length()), CPeptide.sequence()
	for i in range(max(len(CSpool), len(NSpool))):
		
		try:
			l = len(NSpool[i])
			spaces = ' ' * (NS - l)
			try:
				print NSpool[i], spaces, CSpool[i]
			except IndexError:
				print
		except IndexError:
			spaces = ' ' * (1 + NS)
			print spaces, CSpool[i]
	

#def specFilter4_100(sp):
#	retsp = xlh.MSMS()
#	range = sp.mzRange()
#	n = 1
#	while 1:
#		start = range[0] + (n -1) * 100
#		end = range[0] + n * 100
#		retsp.update(sp.spectrumSlice(start, end).mostIntense(4))
#		n += 1
#		if start > range[1]:
#			break
#	retsp.setDescription(sp.description())
#	retsp.setCharge(sp.charge())
#	retsp.setPrecursorMz(sp.precursorMz())
#	return retsp

def mod2string(mod):
	s = ''
	for (pos, mod) in mod.items():
		if mod in FixedModifications:
			continue
		s = s + mod + ':' + str(pos) + ';'
	return s[:-1]

def finneganScan(descr):
#	p = descr.find('precIntensity')
	p = descr.find('FinneganScanNumber')
	return descr[p:]
			

	
def exitOnError(errmsg, exit):
	sys.stderr.write('The following error(s) occurred:\n')
	sys.stderr.write(errmsg + '\n')
	sys.exit(exit)
def usage():
	print sys.argv[0], """-- a tool to validate OMSSA results in a X-link framework
	
Usage:

""", sys.argv[0], """ [-h] -d DATABASE -i INFILE -f FILE [-t TOL] [-u TOL unit] [-e PPM] [-R Results type] [-l AA]
			 [-g g-score] [-p p-value] [-q FDR] [-v] [-o OUTFILE] [-m mask] [-x] [-A] [-c] [-n xlinks]
			 [-r ratio] [-D] [-F fx mod]

 -h:	Print this help and exit
 -d:	Database used for first MS search (not xlink database!) [Default: nr]
 -i:	File containing OMMSA results [only CSV supported]
 -f:	File with spectra [only MGF supported]
 -t:	Tolerance for ion matching [Default 0.5 Da]
 -u:	Tolerance unit [Default Da]
 -e:	ppm error for precursor match [Default 10]
 -R:	OMSSA results format [csv|oms|omx] [not implemented]
 -l:	Filter peptide length [Default: 3 aa]
 -g:	Filter g-scores under this threshold [Default: 0.0]
 -p:	Filter p-values under this threshold [Default: 0.05]
 -q:	Filter q-values (FDR) under this threshold [Default: None]
 -v:	Verbose output 
 -o:	Basename for output files
 -m:	bit mask or value: 5 bit 00000 to mask x-link type
            1   1   1   1   1	
            |   |   |   |   +-> has multiple spectra in interaction
            |   |   |   +-----> has both x-mod position in a spectrum
            |   |   +---------> has different xlinkers
            |   +-------------> has different sequences in xlinks
            +-----------------> has head/tail and tail/head xlinks
 -x:	Include Xl-ions in scoring function [Default: False]
 -A:	Exclude heuristic false matches [Default: False] (Disabled)
 -c:	csv output [Default: False]
 -n:	min. number of xlinks per interaction [Default: 1]
 -r:	ratio for noise reducer [Default: 0.03]
 -D:	Disable score filters (for debug) [Default: False]
 -F:	Fixed Modifications, comma separated, escaped string [Default: 'carbamidomethyl C']
	"""
	sys.exit(1)

OnlyPos = False
AllIons = False
DbIn = 'nr'
ResultsFile = ''
ResultsType = 'csv'
SpectraFile = ''
ShortFilter = 3
Tolerance = 0.5
TolUnit = 'Da'
ppmThreshold = 10
pThreshold = 0.05
gThreshold = 0.0
qThreshold = 1
Verbose = False
Basename = 'default'
bitMask = 0
csvOutput = False
noiseRatio = 0.03
FilterOff = False
FixedModifications = ['carbamidomethyl C']
MinNumberOfSpectra = 1
RankMode = 'gt'

try:
	optlist, args = getopt.getopt(sys.argv[1:], 'hd:i:f:R:t:u:l:e:p:g:q:vo:m:xAcr:DF:n:')
except getopt.GetoptError:
	usage()

if not optlist:
	usage()

for o, a in optlist:
	if o == '-h':
		usage()
	if o == '-d':
		DbIn = a
	if o == '-i':
		ResultsFile = a
	if o == '-f':
		SpectraFile = a
	if o == '-l':
		ShortFilter = int(a)
	if o == '-t':
		Tolerance = float(a)
	if o == '-u':
		TolUnit = a
		if TolUnit not in ['Da', 'ppm']:
			exitOnError(TolUnit + " is not a valid unit, asshole!", 1)
	if o == '-e':
		ppmThreshold = float(a)
	if o == '-p':
		pThreshold = float(a)
	if o == '-g':
		gThreshold = float(a)
	if o == '-q':
		qThreshold = float(a)
	if o == '-R':
		ResultsType = a
		if ResultsType not in ['csv', 'oms', 'omx']:
			exitOnError(ResultsType + " is not a valid format", 1)
		exitOnError("Results Type specification hasn't been implemented yet", 1)
	if o == '-v':
		Verbose = True
	if o == '-o':
		Basename = a
	if o == '-x':
		AllIons = True
	if o == '-A':
		OnlyPos = True
	if o == '-c':
		csvOutput = True
	if o == '-r':
		noiseRatio = float(a)
	if o == '-m':
		mask = int(a)
		if len(a) == 5:
			try:
				bitMask = int(a, 2)
				RankMode = 'bit'
			except ValueError:
				exitOnError('Wrong bitmask, should be 5-bit', 1)
		elif mask < 0 or mask > 31:
			exitOnError('Wrong bitmask, should be 0<x<31', 1)
			bitMask = mask
		else:
			bitMask = mask
			RankMode = 'gt'
	if o == '-D':
		FilterOff = True
	if o == '-F':
		FixedModifications = map(lambda x: x.strip(), a.split(','))
	if o == '-n':
		MinNumberOfSpectra = int(a)
		
		
# do some checks
if not DbIn:
	exitOnError('Missing database name', 1)
if not ResultsFile:
	exitOnError('Missing results file', 1)
if not SpectraFile:
	exitOnError('Missing spectra file', 1)
if ppmThreshold < 0:
	exitOnError('Why you do expect a negative ppm?', 1)
if pThreshold < 0:
	exitOnError('Why you do expect a negative p-value', 1)
	
for mod in FixedModifications:
	if not xlh.Modifications.has_section(mod):
		exitOnError('Modification ' + mod + ' is not in modifications.ini file', 1)
#### End of option parsing section

# Define some variables
# chech which modifications are xlinker

Acclist = {}	# a dictionary of proteins
xlList = xlh.XLinkList()	# a list of XLink xlinks


if Verbose:
	# print operational parameters
	paramOut = "[ Operational Parameters ]\n"
	paramOut += "Database: " + DbIn + '\n'
	paramOut += "OMSSA Results File: " + os.path.abspath(ResultsFile) + '\n'
	paramOut += "Spectra File: " + os.path.abspath(SpectraFile) + '\n'
	paramOut += "Ion Tolerance: " + str(Tolerance) + " " + TolUnit + '\n'
	paramOut += "Precursor Tolerance: " + str(ppmThreshold) + " ppm\n"
	paramOut += "Min. Peptide Length: " + str(ShortFilter) + '\n'
	paramOut += "G-Score Threshold: " + str(gThreshold) + '\n'
	paramOut += "P-Value Threshold: " + str(pThreshold) + '\n'
	paramOut += "FDR Threshold: " + str(qThreshold) + '\n'
	paramOut += "Mask: " + str(bitMask) + '\n'
	paramOut += "Noise Ratio: " + str(noiseRatio) + '\n'
	paramOut += "Min. number of spectra: " + str(MinNumberOfSpectra) + '\n'
	paramOut += "Include Xl-ions: " + str(AllIons) + '\n\n'
#	paramOut += "\nPrint: " + ResultsFile + '\n'
	sys.stderr.write(paramOut)
	

sys.stderr.write('Process started on ' + ctime() + '\n')
now = time()

sys.stderr.write('STEP 1: Read results and filter matches\n')

sys.stderr.write('Reading results file... ')
sys.stderr.flush()

raw = xlh.OMSSAResults()
raw.parseOMSSACsv(ResultsFile)

#if Verbose:
ok = 'OK [ ' + str(raw.length()) + ' matches ]\n'
#else:
#	ok = 'OK\n'

sys.stderr.write(ok)

sys.stderr.write('Reading spectra file... ')
sys.stderr.flush()
spectra = xlh.MSMSAnalysis()
# read as mgf
spectra.parseMgfFile(SpectraFile)
#spectra.parseMzXML(SpectraFile)
	
#if Verbose:
ok = 'OK [ ' + str(spectra.length()) + ' spectra ]\n'
#else:
#	ok = 'OK\n'
sys.stderr.write(ok)

sys.stderr.write('Reducing noise... ')
sys.stderr.flush()
spectra.reduceNoise(noiseRatio)
ok = 'OK\n'
sys.stderr.write(ok)


sys.stderr.write('Checking for data consistency... ')
sys.stderr.flush()
filtered = raw.checkPeptideConsistency(ShortFilter)
			
#if Verbose:
ok = 'OK [ ' + str(filtered.length()) + ' matches ]\n'
#else:
#	ok = 'OK\n'
sys.stderr.write(ok)

del raw

if filtered.length() == 0:
	del spectra
	sys.stderr.write('*** No valid matches for this search! ***\n')
	sys.exit(0)


if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 1\n')
	now = time()

sys.stderr.write('STEP 2: Validate remaining matches and protein mapping\n')


sys.stderr.write('Reading protein sequences... ')
sys.stderr.flush()
Acclist = filtered.proteins(DbIn)

#if Verbose:
ok = 'OK [ ' + str(len(Acclist)) + ' sequences ]\n'
#else:
#	ok = 'OK\n'
sys.stderr.write(ok)

sys.stderr.write('Guessing x-link sites... ')
sys.stderr.flush()

xlList = filtered.guessXLinks(spectra, Tolerance, TolUnit)
xlList.setFixedModifications(FixedModifications)
###3			
#			if AllIons:
#				sp = specFilter4_100(spectra.descriptionSearch(m.description()).spectra()[0])
#			else:
#				sp = spectra.descriptionSearch(m.description()).spectra()[0]
#			
####3		

#filter on error		
#if not FilterOff:
xlList = xlList.errorFilter(ppmThreshold)

# calculate p-values
xlList.calculatePValues(Tolerance, TolUnit, AllIons)

#filter on p-value
if not FilterOff:
	xlList = xlList.pValueFilter(pThreshold)
	xlList.correctByFDR(qThreshold)

#if Verbose:
ok = 'OK [ ' + str(xlList.length()) + ' xlinks ]\n'
#else:
#	ok = 'OK\n'
sys.stderr.write(ok)

if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 2\n')
	now = time()

if xlList.length() == 0:
	del spectra
	del xlList
	sys.stderr.write('*** No x-links for this search! ***\n')
	sys.exit(0)


sys.stderr.write('STEP 3: Build xlink table\n')

sys.stderr.write('Collapsing xlinks... ')
sys.stderr.flush()

#pdb.set_trace()


xlList.guessInteractions()

# do some filtering...
if not FilterOff:
	xlList = xlList.gScoreFilter(gThreshold)
	xlList = xlList.rankFilter(bitMask, RankMode)
	xlList = xlList.specLengthFilter(MinNumberOfSpectra)
#if OnlyPos:
#	xlList = filter(lambda x: x.rank() != 0 and x.gScore() != '-Inf', xlList._eList)


#if Verbose:
ok = 'OK [ ' + str(xlList.eLength()) + ' interactions ]\n' 
#else:
#	ok = 'OK\n'
sys.stderr.write(ok)

#######
#### trying to get some true positives... we take for each rank the first quartile as threshold
#### for rank higher than 9 the threshold should not be less than -10
###
###rankDict = {}
###thrDict = {}
###for r in dict.fromkeys(map(lambda x: x.rank(), xlList.events())):
###	rankDict[r] = map(lambda x: x.gScore(), filter(lambda x: x.rank() == r, xlList.events()))
###	rankDict[r].sort()
###	fi = len(rankDict[r]) / 4.0
###	if fi <=1:
###		thrIndex = 0
###	elif abs(fi * 2 - int(fi) * 2) > 1:
###		# Haven't found how to have a nice integer casting...
###		thrIndex = int(ceil(fi)) - 1
###	else:
###		thrIndex = int(fi) - 1
###	
###	if r >= 8 and rankDict[r][thrIndex] > -10:
###		thrDict[r] = -10.0
###	elif r >= 2 and rankDict[r][thrIndex] > -8:
###		thrDict[r] = -8.0
###	elif r >= 0 and rankDict[r][thrIndex] > -5:
###		thrDict[r] = -5.0
###	else:
###		thrDict[r] = rankDict[r][thrIndex]
###	
###	sys.stderr.write('Threshold for rank ' + str(r) + ': ' + str(thrDict[r]) + '\n') 
#######


# Try to build a rough CDF for g-scores, so that we can set an upper limit, say at 0.95 cum
gScoreList = map(lambda x: x.gScore(), xlList.interactions())
#gScoreList.sort()
#for g in gScoreList:
#	eventsBelow = filter(lambda x: x.gScore() <= g, xlList.events())
#	n = sum(map(lambda x: x.length(), eventsBelow))
#	if n >= 0.9 * xlList.length():
#		sys.stderr.write('g-score threshold (based on CDF): ' + str(g) + '\n')
#		break



sys.stderr.write('*** Warning!!! This is just a raw list ***\n')

xlList.sort(key = 'gScore')


if csvOutput:
	print 'L-Protein, R-Protein, Start, End, Rank, G-Score, L-Peptide, modstring, R-Peptide, modstring, p-value, error, m/z, charge, number, description'
	for x in xlList.interactions():
		for m in x.Xlinks():
			ppmerror = (m.source().precursorMass() - m.source().precursorTheoMass()) / m.source().precursorMass() * 1000000
			print ','.join(map(lambda x: str(x), [x.NProtein().accession(), x.CProtein().accession(), x.start(), x.end(), x.rank(), x.gScore(), 
			m.NPeptide().sequence(), mod2string(m.NPeptide().modifications()), 
			m.CPeptide().sequence(), mod2string(m.CPeptide().modifications()), 
			m.pValue(), ppmerror, m.source().precursorMz(), m.source().charge(), 
			m.source().spectrumNumber() + 1, finneganScan(m.spectrum().description())]))
else:
	nev = 1
	for x in xlList.interactions():
		print '[ Interaction', nev, ']'
		print 'Coordinates:', x.NProtein().accession() + ':' + str(x.start()), x.CProtein().accession() + ':' + str(x.end())
		print 'G-Score:', x.gScore(), ' [', correct(percentileofscore(gScoreList, x.gScore())), ']'
#		print 'Combined p-value:', combinedPV(x)
		print 'Rank:', x.rank()
		print 'N. of xlinks:', x.length()	
		print 'N. of spectra:', x.numberOfSpectra()
		nxl = 1
		if Verbose:
			print 'Xlinks'
			for m in x.xLinks():
				ppmerror = (m.source().precursorMass() - m.source().precursorTheoMass()) / m.source().precursorMass() * 1000000
				print '[', nxl, ']'
				print '\tLeft Peptide:', m.NPeptide().sequence(), '[', mod2string(m.NPeptide().modifications()), ']'
				print '\tRight Peptide:', m.CPeptide().sequence(), '[', mod2string(m.CPeptide().modifications()), ']'
				print '\tp-value:', m.pValue()
				print '\tSpectrum Number:', m.source().spectrumNumber() + 1, '[',finneganScan(m.spectrum().description()),']'
				print '\tCharge:', m.source().charge(), '\tm/z:', m.source().precursorMz(), '\tError (ppm):', ppmerror
				nxl += 1
				if Verbose:
					printMatchingIons(m, Tolerance, TolUnit)
		else:
			#print at least some sequence....
			print 'Sequences:', x.NProtein().subsequence(center = x.start()).sequence(), x.CProtein().subsequence(center = x.end()).sequence()
		nev += 1
		print
if Verbose:
	sys.stderr.write('Took ' + str(int(time() - now)) + ' s to complete step 3\n')

sys.stderr.write('Process ended on ' + ctime() + '\n')
if FilterOff:
	sys.stderr.write('FDR: ' + str(xlList.fdr()) + '\n')


### Print statistics
print "#######################"
print "Number of Interactions", xlList.eLength()
print "gScore stats:", np.mean(gScoreList), np.std(gScoreList), np.median(gScoreList)
nXlinkList = map(lambda x: x.length(), xlList.interactions())
print "Xlink stats:", np.mean(nXlinkList), np.std(nXlinkList), np.median(nXlinkList)
rankList = map(lambda x: x.rank(), xlList.interactions())
print "Rank stats:", np.mean(rankList), np.std(rankList), np.median(rankList)
print "#######################"
del xlList
